use crate::app::AppState;
use crate::config;
use crate::error::{ApiError, ApiResult};
use axum::debug_handler;
use axum::extract::{FromRef, FromRequestParts};
use axum::http::request::Parts;
use jsonwebtoken::jwk::{Jwk, JwkSet};
use jsonwebtoken::{decode, decode_header, DecodingKey, Header, Validation};
use once_cell::sync::Lazy;
use reqwest::header;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use url::Url;

static JWKS_CACHE: Lazy<RwLock<HashMap<String, Jwk>>> = Lazy::new(|| RwLock::new(HashMap::new()));
#[derive(Debug, Clone)]
pub struct OidcClient {
    pub issuer_uri: Url,
    pub audience: String,
    pub jwks_uri: Url,
}


impl OidcClient {
    pub async fn new() -> ApiResult<Arc<Self>> {
        let auth_config = config::get().auth();
        let mut issuer_uri = String::from(auth_config.issuer_uri());
        issuer_uri.push('/');
        let issuer_uri = Url::parse(issuer_uri.as_str())?;
        let well_known_url = issuer_uri.join(".well-known/openid-configuration")?;
        let discovery_doc: serde_json::Value = reqwest::get(well_known_url)
            .await?
            .json()
            .await?;
        let jwks_uri_str = discovery_doc["jwks_uri"]
            .as_str()
            .ok_or_else(|| ApiError::AuthServerError(String::from("jwks_uri not found in discovery document")))?;

        let jwks_uri = Url::parse(jwks_uri_str)?;
        tracing::info!("jwks_uri: {:?}", jwks_uri);
        let client = Arc::new(Self {
            issuer_uri: Url::parse(auth_config.issuer_uri())?,
            audience: String::from(auth_config.audience()),
            jwks_uri,
        });
        client.refresh_jwks().await?;
        Ok(client)
    }

    pub async fn refresh_jwks(&self) -> ApiResult<()> {
        let jwks_set: JwkSet = reqwest::get(self.jwks_uri.clone())
            .await?
            .json()
            .await?;
        // 拿到写锁
        tracing::info!("jwks_set: {:#?}", jwks_set);
        let mut guard_cache_write = JWKS_CACHE.write().await;
        guard_cache_write.clear();
        for jwk in jwks_set.keys {
            if let Some(kid) = jwk.common.key_id.clone() {
                guard_cache_write.insert(kid, jwk);
            }
        }
        tracing::info!("JWKS refreshed successfully. Cached {} keys.", guard_cache_write.len());
        Ok(())
    }

    async fn get_jwk(&self, header: &Header) -> Option<Jwk> {
        let cache_read = JWKS_CACHE.read().await;
        let kid = header.kid.as_ref()?;
        cache_read.get(kid).cloned()
    }

    pub async fn validate_token(&self, token: &str) -> ApiResult<Claims> {
        //1. 从token拿出header
        let header = decode_header(token).map_err(|_| {
            ApiError::AuthServerError("Decode token failed".to_string())
        })?;

        // 2. 从根据header的kid从Cache拿到jwk
        let jwk = match self.get_jwk(&header).await {
            Some(jwk) => jwk,
            None => {
                //不存在jwk，尝试刷新
                self.refresh_jwks().await.map_err(|_| {
                    ApiError::AuthServerError("Refresh token error".to_string())
                })?;
                let Some(jwk_after_refresh) = self.get_jwk(&header).await else {
                    return Err(ApiError::AuthServerError("Refresh token ineffective".to_string()));
                };
                jwk_after_refresh
            }
        };


        //3. 创建DecodingKey & Validation
        let decoding_key = DecodingKey::from_jwk(&jwk).map_err(|_| {
            ApiError::AuthServerError("Failed to create decoding key from JWK".to_string())
        })?;
        let mut validation = Validation::new(header.alg);
        validation.set_issuer(&[self.issuer_uri.to_string()]);
        validation.set_audience(&[self.audience.clone()]);
        validation.validate_exp = true;
        // validation.validate_nbf = true;
        // validation.set_required_spec_claims(&["exp", "iat", "iss", "aud", "sub", "typ"]);

        //4. token验签
        let token_data = decode(token, &decoding_key, &validation).map_err(|e| {
            ApiError::UnAuthenticatedError(format!("Token Valid Failed: {:?}", e))
        })?;
        Ok(token_data.claims)
    }
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Claims {
    pub iss: String,
    pub sub: String,
    pub aud: serde_json::Value,
    pub exp: usize,
    pub nonce: Option<String>,
    pub preferred_username: Option<String>,
    pub name: Option<String>,
    pub email: Option<String>,
    pub iat: usize,
}


// Keycloak 提供的 Realm Access 声明结构
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct RealmAccess {
    pub roles: Vec<String>,
}

// Keycloak 提供的 Resource Access 声明结构
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ResourceAccess {
    pub roles: Vec<String>,
}

#[derive(Debug, Clone)]
pub struct OidcClaims(pub Claims);

// 从请求头中提取OidcClaims 实现FromRequestParts
impl<S> FromRequestParts<S> for OidcClaims
where
    S: Send + Sync,
    AppState: FromRef<S>,
{
    type Rejection = ApiError;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let auth_header = parts.headers
            .get(header::AUTHORIZATION)
            .and_then(|v| v.to_str().ok());
        let token = match auth_header {
            Some(authentication) if authentication.starts_with("Bearer ") => authentication.strip_prefix("Bearer ").unwrap().to_string(),
            _ => return Err(ApiError::UnAuthenticatedError("Bearer token invalid".to_string())),
        };
        let app_state = AppState::from_ref(state);
        let claims = app_state.oidc_client.validate_token(&token).await?;
        parts.extensions.insert(claims.clone());
        Ok(OidcClaims(claims))
    }
}

